import os
import re
import numpy as np
import shutil

def CheckFile(ifile):
    if not os.path.isfile(ifile):
        print("Can not find file %s" % ifile)
        return False
    return True

def get_block(lines, keyword, skip=0, nlines=None):
    ret = []
    found = False
    if not nlines:
        nlines = 1e6
    for idx, ii in enumerate(lines):
        if keyword in ii:
            found = True
            blk_idx = idx + 1 + skip
            line_idx = 0
            while len(re.split("\s+", lines[blk_idx])) == 0:
                blk_idx += 1
            while line_idx < nlines and blk_idx != len(lines):
                if len(re.split("\s+", lines[blk_idx])) == 0 or lines[blk_idx] == "":
                    blk_idx += 1
                    continue
                if "coxmol: Writing XMOL coordinates into file siesta.xyz" in lines[blk_idx]:
                    break
                ret.append(lines[blk_idx])
                blk_idx += 1
                line_idx += 1    
            break
    if not found:
        return None
    return ret


def get_geometry_in(fname, inlines):
    geometry_path_in = os.path.join(fname, "output")
    for line in inlines:
        if "fdf_file" in line and "fdf_file" == line.split()[0]:
            atom_file = line.split()[1]
            geometry_path_in = os.path.join(fname, atom_file)
            break
    return geometry_path_in


def get_output_cell(geometry_inlines):
    cell_lines = get_block(geometry_inlines,"outcell: Unit cell vectors (Ang):",skip=0,nlines=3)
    cell = []
    for line in cell_lines:
        parts = line.split()
        cell.append([float(coord)for coord in parts[:3]])
    cell = np.array(cell)

    return cell

def get_output_coords(geometry_inlines, inlines=None):
    atom_species = []
    coords = []

    atom_specie = [line for line in geometry_inlines if "NumberOfSpecies" in line][0]
    atom_species.append(int(atom_specie.split()[-1]))

    coords_line = get_block(geometry_inlines, "outcoor: Relaxed atomic coordinates (Ang):",nlines=1e6)
    for line in coords_line:
        parts = line.split()
        coords.append([float(coord) for coord in parts[:3]])
    coords = np.array(coords)
    np.set_printoptions(suppress=True)
    return atom_species,coords

def get_siesta_output(input_file,INPUT=None):
    with open(input_file, "r") as fp:
        geometry_inlines = fp.read().split("\n")
    for iline, line in enumerate(geometry_inlines):
        if line.split() == [] or len(line) == 0:
            del geometry_inlines[iline]
    geometry_inlines.append("")
    cell = get_output_cell(geometry_inlines)
    species,coords = get_output_coords(geometry_inlines)
    
    data = {}
    data["atom_species"] = species
    data["cells"] = cell
    data["coords"] = coords
    return data

def make_siesta_out_fdf(sys_data,output_file):
    new_data = []
    new_species = sys_data["atom_species"]
    new_cell = sys_data["cells"]
    new_coords = sys_data["coords"]
    with open(output_file, "r") as file:
        olddata = file.readlines()
        delete_mode = False
        for line in olddata :
            if "%block LatticeVectors" in line:
                delete_mode = True
                new_data.append(line)
                continue
            elif "%endblock LatticeVectors" in line:
                if len(new_cell) > 0:
                    new_line = ""
    
                    for i, c in enumerate(new_cell):
                        formatted_c = ["{:7.6f}".format(coord) for coord in c]
                        if i == 0:
                            new_line += "  " + "   ".join(formatted_c) + "\n"
                        else:
                            new_line += "  " + "   ".join(formatted_c) + "\n"
                    new_data.append(new_line)
                delete_mode = False
                new_data.append(line)
                continue
            elif "%block AtomicCoordinatesAndAtomicSpecies" in line:
                delete_mode = True
                new_data.append(line)
                continue
            elif "%endblock AtomicCoordinatesAndAtomicSpecies" in line:
                if len(new_coords) > 0:
                    new_coord = ""
                    for c in new_coords:
                        formatted_c = [f"{coord:.6f}" for coord in c]
                        species = str(new_species[0])
                        formatted_c.append(species)
                        new_coord += " ".join(formatted_c) + "\n"
                    new_coord = new_coord.strip() + "\n"
                    new_data.append(new_coord)
                delete_mode = False
                new_data.append(line)
                continue
            if delete_mode:
                continue
            new_data.append(line)
    with open(output_file,"w") as file:
        file.writelines(new_data)
